import os
from CalTorque import CalTorque
import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
from matplotlib import dates
import datetime


def CompareSources(data_files, bins, max_d):
    src = {}
    file = open('run_history.txt')
    lines = file.readlines()
    
    d = datetime.datetime.now()
    q = str(d.strftime('%Y%m%d'))
    
    for line in lines:
        source = line.split(', ')[2].split(':')[1] + line.split(', ')[2].split(':')[0]
        if source not in src.keys():
            src[source] = []
    
    for file in data_files:
        data = CalTorque(file)
        #if data['station'] != 'S5': continue
        #if data['position'][0] > 50: continue
        #if max(data['pos_offsetfix']) < 350: continue
        if int(data['run']) <= 20130222: # do not consider anything taken with the old telescope
          continue
        print file
        source = data['source']
        if source in src.keys():
            src[source].append([data['pos_offsetfix'], data['torque_inoz']])
    
    countRa = len(src['weakRa-226'])
    countCs = len(src['weakCs-137'])
    countCo = len(src['weakCo-60'])
    countTh = len(src['weakTh-228'])
    
    for s in src:
        #print len(src[s])
        x_inputs = numpy.zeros(bins)
        X = numpy.linspace(0,max_d,bins)
        Y = numpy.zeros(bins)
        for i, k in enumerate(src[s]):
            # k[0] = x, k[1] = y
            for j, x in enumerate(X):
                pos_torque = []
                x0 = x-(max_d/bins)/2.
                x1 = x+(max_d/bins)/2.
                for i, p in enumerate(k[0]):
                    if p > x0 and p <= x1:
                        pos_torque.append(k[1][i])
                if len(pos_torque)>0:
                    if max(pos_torque) == 0.: continue
                    Y[j] += max(pos_torque)
                    x_inputs[j]+=1.
        
        X_nonull = []
        Y_nonull = []
        x_inputs_nonull = []
        
        # gets rid of empty bins
        for g, h in enumerate(x_inputs):
            if h == 0: continue
            X_nonull.append(X[g])
            Y_nonull.append(Y[g])
            x_inputs_nonull.append(x_inputs[g])
        
        for l, n in enumerate(x_inputs_nonull):
            if n==0: continue
            Y_nonull[l] = Y_nonull[l]/n
        src[s] = [X_nonull, Y_nonull]

    # plot source averages
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax = fig.add_axes(box)
    ax.set_xlabel('Position (cm)')
    ax.set_ylabel('Torque (in.oz)')
    maxx = max(data['pos_offsetfix'])
    miny = 0
    
    ## Add station locations
    if maxx >= 328.9:
        ax.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## Add bend locations
    if maxx >= 47.8:
        ax.axvspan(47.8, 93.5, color = 'g', alpha = 0.2)
        ax.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax.axvspan(136.4, 144.6, color = 'g', alpha = 0.2)
        ax.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax.axvspan(196.5, 218.0, color = 'g', alpha = 0.2)
        ax.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax.axvspan(228.3, 241.5, color = 'g', alpha = 0.2)
    if maxx >= 274.6:
        ax.axvspan(274.6, 284.6, color = 'g', alpha = 0.2)
        ax.annotate('Bend A5', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax.axvspan(309.0, 318.9, color = 'g', alpha = 0.2)
        ax.annotate('Bend A6', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax.axvspan(358.8, 370.8, color = 'g', alpha = 0.2)
        ax.annotate('Bend A7', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax.axvspan(414.5, 426.5, color = 'g', alpha = 0.2)
        ax.annotate('Bend A8', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax.axvspan(445.3, 455.3, color = 'g', alpha = 0.2)
        ax.annotate('Bend A9', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax.axvspan(467.5, 479.5, color = 'g', alpha = 0.2)
        ax.annotate('Bend A10', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax.axvspan(523.8, 535.8, color = 'g', alpha = 0.2)
        ax.annotate('Bend A11', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax.axvspan(571.6, 583.6, color = 'g', alpha = 0.2)
        ax.annotate('Bend A12', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax.axvspan(627.9, 639.9, color = 'g', alpha = 0.2)
        ax.annotate('Bend A13', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax.axvspan(652.7, 660.7, color = 'g', alpha = 0.2)
        ax.annotate('Bend A14', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax.axvspan(678.5, 681.1, color = 'g', alpha = 0.2)
        ax.annotate('Bend A15', xy = (678, miny), color = 'g')
    
    colors = "bgrcmy"
    color_index = 0
    for sources in src:
        if src[sources][1] == []: continue
        if sources == 'weakTh-228':
            ax.plot(src[sources][0], src[sources][1], '-', linewidth = 2, c='k', label = sources[4:])
        else:
            ax.plot(src[sources][0], src[sources][1], '-', linewidth = 2, c=colors[color_index], label = sources[4:])
            color_index += 1
    h, l = ax.get_legend_handles_labels()
    ax.grid()
    PLT.title('Average Torque for Different Sources')
    PLT.legend(h, l, 'upper left')
    PLT.savefig('Sources_averages_Updated-' + q + '.png')
    
    print "Number of Radium runs: ", countRa
    print "Number of Cesium runs: ", countCs
    print "Number of Cobalt runs: ", countCo
    print "Number of Thorium runs: ", countTh